package com.aiwiown.snackmq.network.handler;

import com.aiwiown.snackmq.common.message.Message;
import com.aiwiown.snackmq.common.message.MessageType;
import io.netty.channel.Channel;
import io.netty.channel.ChannelFuture;
import io.netty.channel.ChannelHandler;
import io.netty.channel.ChannelHandlerContext;
import io.netty.channel.SimpleChannelInboundHandler;
import lombok.extern.slf4j.Slf4j;

import java.io.IOException;
import java.util.Map;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.atomic.AtomicLong;

@Slf4j
@ChannelHandler.Sharable
public class ClientHandler extends SimpleChannelInboundHandler<Message> {

    private final Map<Long, CompletableFuture<Message>> pendingRequests = new ConcurrentHashMap<>();
    private final AtomicLong messageIdGenerator = new AtomicLong(1); // 从 1 开始，避免 0

    /**
     * 处理接收到的消息。
     *
     * @param ctx 上下文对象，包含与通道相关的状态信息。
     * @param msg 接收到的消息对象。
     */
    @Override
    protected void channelRead0(ChannelHandlerContext ctx, Message msg) {
        // 如果消息类型是响应类型，则处理响应消息
        if (msg.getType() == MessageType.RESPONSE) {
            log.debug("Received response: {}", msg);
            CompletableFuture<Message> future = pendingRequests.remove(msg.getMessageId());
            if (future != null) {
                // 将响应消息传递给对应的 CompletableFuture
                future.complete(msg);
            } else {
                // 如果找不到对应的请求 ID，记录警告日志
                log.warn("Received response for unknown messageId: {}", msg.getMessageId());
            }
        } else {
            // 如果接收到的不是响应类型的消息，记录警告日志并继续传递消息
            log.warn("Client received a non-response message type: {}", msg.getType());
            ctx.fireChannelRead(msg);
        }
    }

    /**
     * 发送请求消息并返回一个 CompletableFuture，用于异步接收响应。
     *
     * @param message 请求消息对象。
     * @param channel 用于发送消息的通道。
     * @return 一个 CompletableFuture，当接收到响应时完成。
     */
    public CompletableFuture<Message> sendRequest(Message message, Channel channel) {
        // 参数验证
        if (message == null) {
            throw new IllegalArgumentException("Message cannot be null");
        }
        if (channel == null) {
            throw new IllegalArgumentException("Channel cannot be null");
        }
        
        // 生成唯一的请求 ID
        final long messageId = messageIdGenerator.getAndIncrement();
        message.setMessageId(messageId);

        // 创建 CompletableFuture 并将其与请求 ID 关联
        CompletableFuture<Message> future = new CompletableFuture<>();
        pendingRequests.put(messageId, future);

        // 发送消息并通过监听器处理发送失败的情况
        ChannelFuture channelFuture = channel.writeAndFlush(message);
        if (channelFuture != null) {
            channelFuture.addListener(f -> {
                if (!f.isSuccess()) {
                    // 如果发送失败，移除请求并标记 CompletableFuture 为异常完成
                    pendingRequests.remove(messageId);
                    future.completeExceptionally(f.cause());
                }
            });
        } else {
            // 如果 writeAndFlush 返回 null，标记为异常
            pendingRequests.remove(messageId);
            future.completeExceptionally(new IllegalStateException("Channel writeAndFlush returned null"));
        }

        return future;
    }

    /**
     * 处理通道中捕获的异常。
     *
     * @param ctx 上下文对象，包含与通道相关的状态信息。
     * @param cause 异常对象。
     */
    @Override
    public void exceptionCaught(ChannelHandlerContext ctx, Throwable cause) {
        log.error("Exception caught in client handler", cause);
        // 清理所有待处理的请求，并标记为异常完成
        pendingRequests.values().forEach(future -> future.completeExceptionally(cause));
        pendingRequests.clear();
        ctx.close();
    }

    /**
     * 处理通道变为非活动状态的情况。
     *
     * @param ctx 上下文对象，包含与通道相关的状态信息。
     * @throws Exception 如果在处理过程中发生异常。
     */
    @Override
    public void channelInactive(ChannelHandlerContext ctx) throws Exception {
        log.warn("Channel became inactive: {}", ctx.channel().remoteAddress());
        // 当连接断开时，创建一个 IOException 并标记所有待处理的请求为异常完成
        IOException ex = new IOException("Connection to " + ctx.channel().remoteAddress() + " is closed.");
        pendingRequests.values().forEach(future -> future.completeExceptionally(ex));
        pendingRequests.clear();
        super.channelInactive(ctx);
    }
}
